#ifndef MULTIGRIDTWODIM_HPP
#define MULTIGRIDTWODIM_HPP

#include "jacobi.hpp"
#include "option2D.hpp"
#include "norm.hpp"
#include "point.hpp"
#include <math.h>

using namespace std;
using namespace Eigen;
/*****************************************class function*****************************************/
// class Function
// {
// public:
//     Function() {}
//     double operator()(point _p)
//     {
//         return pow(_p.x,2) + pow(_p.y,2);
//     }

//     double f_x(point _p)
//     {
//         return 2 * _p.x;
//     }

//     double f_y(point _p)
//     {
//         return 2 * _p.y;
//     }
//     /**********************************负拉普拉斯算子**********************************/
//     double laplace_f(double _x, double _y)
//     {
//         double fxx = 2.0;
//         double fyy = 2.0;
//         return -fxx - fyy;
//     }

//     double laplace_f(point _p)
//     {
//         return laplace_f(_p.x, _p.y);
//     }
// };
// class Function
// {
// public:
//     Function() {}
//     double operator()(point _p)
//     {
//         return exp(_p.y + sin(_p.x));
//     }

//     double f_x(point _p)
//     {
//         return cos(_p.x) * exp(_p.y + sin(_p.x));
//     }

//     double f_y(point _p)
//     {
//         return exp(_p.y + sin(_p.x));
//     }
//     /**********************************负拉普拉斯算子**********************************/
//     double laplace_f(double _x, double _y)
//     {
//         double fxx = (-sin(_x) + pow(cos(_x), 2)) * exp(_y + sin(_x));
//         double fyy = exp(_y + sin(_x));
//         return -fxx - fyy;
//     }

//     double laplace_f(point _p)
//     {
//         return laplace_f(_p.x, _p.y);
//     }
// };

class Function
{
public:
    Function() {}
    double operator()(point _p)
    {
        return sin(_p.x + _p.y);
    }

    double f_x(point _p)
    {
        return cos(_p.x + _p.y);
    }

    double f_y(point _p)
    {
        return cos(_p.x + _p.y);
    }
    /**********************************负拉普拉斯算子**********************************/
    double laplace_f(double _x, double _y)
    {
        double fxx = -sin(_x + _y);
        double fyy = -sin(_x + _y);
        return -fxx - fyy;
    }

    double laplace_f(point _p)
    {
        return laplace_f(_p.x, _p.y);
    }
};
/*****************************************二维多重网格*****************************************/
class MultigridTwoDim
{
private:
    int n; // n个interval
    double h;
    vector<vector<point>> GridTable;
    Function func;
    MatrixXd A, Ah;
    VectorXd f, sol,
        vec_f, vec_sol;

public:
    /****************************************************************************Set Grid**********************************************************************************/
    vector<vector<point>> set_grid(int _n);

    /****************************************************************************Set Matrix**********************************************************************************/
    void set_matrix(BoundCondiType bound, int _n); // 分成n个intervals, 32 16 8....

    /****************************************************************************V-cycle**********************************************************************************/
    VectorXd V_cycle(double w, VectorXd v, VectorXd _f, int nu1, int nu2,
                     BoundCondiType bound, Operators::Restriction rest, Operators::Interpolation inter)
    {
        int _n = sqrt(v.size()) - 1;
        set_matrix(bound, _n);
        /***********************VC-1:Relax nu1 times***********************/
        iteration I(Ah, _f, w);
        v = I.weighted_jacobi_iteration(v, nu1);
        /***********************GridSize Judgement***********************/
        if (_n != 2) // not Coarest Grid
        {
            int m = _n / 2 + 1;
            VectorXd v2 = VectorXd::Zero(pow(m, 2));
            /***********************Restriction***********************/
            VectorXd r = _f - Ah * v;
            VectorXd tmp = RestrictionMatrix(rest, _n) * r;
            v2 = V_cycle(w, v2, tmp, nu1, nu2, bound, rest, inter);

            /***********************Interpolation***********************/
            v = InterpolationMatrix(inter, m - 1) * v2 + v;
        }
        /***********************Relax nu1 times***********************/
        // is coarsest grid
        v = I.weighted_jacobi_iteration(v, nu2);
        return v;
    }

    /****************************************************************************FMG**********************************************************************************/
    VectorXd FMG(double w, VectorXd _f, int nu1, int nu2,
                 BoundCondiType bound, Operators::Restriction rest, Operators::Interpolation inter)
    {
        VectorXd vtmp;
        int _n = sqrt(_f.size()) - 1;
        if (_n != 2)
        {
            int m = _n / 2 + 1;
            VectorXd ftmp = RestrictionMatrix(rest, _n) * _f;
            vtmp = FMG(w, ftmp, nu1, nu2, bound, rest, inter);
            vtmp = InterpolationMatrix(inter, m - 1) * vtmp;
        }
        else
        {
            vtmp = VectorXd::Zero(_f.size());
        }
        vtmp = V_cycle(w, vtmp, _f, nu1, nu2, bound, rest, inter);
        return vtmp;
    }

    void Possion_BVP_Multigrid(int size, MultigridMethod method, double w, int nu1, int nu2,
                               BoundCondiType bound, Operators::Restriction rest, Operators::Interpolation inter,
                               int max_iter, double epsi);
    void display()
    {
        cout << "x = [";
        for (int i = 0; i <= n; i++)
        {
            for (int j = 0; j <= n; j++)
                cout << GridTable[i][j].x << ",";
        }
        cout << "];" << endl;

        cout << "y = [";
        for (int i = 0; i <= n; i++)
        {
            for (int j = 0; j <= n; j++)
                cout << GridTable[i][j].y << ",";
        }
        cout << "];" << endl;

        cout << "z = [";
        int cnt = 0;
        for (int i = 0; i < n + 1; i++)
        {
            for (int j = 0; j <= n; j++)
            {
                cout << sol[cnt] << ",";
                cnt++;
            }
        }
        cout << "];" << endl;

        cout << "scatter3(x,y,z, \".\");" << endl;
    }
};

void MultigridTwoDim::Possion_BVP_Multigrid(int size, MultigridMethod method, double w, int nu1, int nu2,
                                            BoundCondiType bound, Operators::Restriction rest, Operators::Interpolation inter,
                                            int max_iter, double epsi)
{
    GridTable = set_grid(size);
    set_matrix(bound, size);
    f = vec_f;
    n = size;
    sol = vec_sol;
    VectorXd real = VectorXd::Zero(sol.size());
    int cnt0 = 0;
    for(int i = 0; i <= n; i++)
    {
        for(int j = 0; j <= n; j++)
        {
            real(cnt0) = func(GridTable[i][j]);
            cnt0++;
        }
    }
    A = Ah;
    Norm Err;
    VectorXd err;
    double derr, Derr;
    if (method == MultigridMethod::Vcycle)
    {
        int cnt = 0;
        for (int i = 1; i <= max_iter; i++)
        {
            /************多重网格法*************/
            sol = V_cycle(w, sol, f, nu1, nu2, bound, rest, inter);
            /************精确解***********/
            err = sol - A.lu().solve(f);
            derr = Err.inftynorm(err);
            Derr = Err.inftynorm(sol - real);
            cout << "V-cycle: " << i << "th: "
                 << "解析解误差： " << Derr << " 离散解误差： " << derr << endl;
            cnt++;
            if (derr < epsi)
                break;
        }
        cout << "err = " << derr << " < epsi: " << cnt << "iteration" << endl;
    }
    else if (method == MultigridMethod::FMG)
    {
        /************多重网格法*************/
        sol = FMG(w, f, nu1, nu2, bound, rest, inter);
        /************精确解*************/
        err = sol - A.lu().solve(f);
        derr = Err.inftynorm(err);
        cout << "FMG: " << 1 << "th:" << derr << endl;
    }
}

vector<vector<point>> MultigridTwoDim::set_grid(int _n)
{
    h = 1.0 / _n;
    vector<vector<point>> p_list;
    /***********************set GridPoint***********************/
    for (int i = 0; i < _n + 1; i++)
    {
        vector<point> vtmp;
        for (int j = 0; j < _n + 1; j++)
        {
            point ptmp(j * h, i * h);
            vtmp.push_back(ptmp);
        }
        p_list.push_back(vtmp);
    }
    return p_list;
}

void MultigridTwoDim::set_matrix(BoundCondiType bound, int _n) // 分成n个intervals, 32 16 8....
{
    double _h = 1.0 / _n;
    int dimention = pow(_n + 1, 2);
    vector<vector<point>> gridtable = set_grid(_n);
    Ah = MatrixXd::Zero(dimention, dimention);
    vec_sol = VectorXd::Zero(dimention);
    vec_f = VectorXd::Zero(dimention);
    int cnt = 0;
    /***********************Dirichlet Boundary***********************/
    if (bound == BoundCondiType::Dirichlet)
    {
        for (int i = 0; i <= _n; i++)
        {
            for (int j = 0; j <= _n; j++)
            {
                if (gridtable[i][j].insquare != 0)
                {
                    Ah(cnt, cnt) = 1;
                    vec_f(cnt) = func(gridtable[i][j]);
                }
                else
                {
                    Ah(cnt, cnt - 1) = -1.0 / pow(_h, 2);
                    Ah(cnt, cnt + 1) = -1.0 / pow(_h, 2);
                    Ah(cnt, cnt - _n - 1) = -1.0 / pow(_h, 2);
                    Ah(cnt, cnt + _n + 1) = -1.0 / pow(_h, 2);
                    Ah(cnt, cnt) = 4.0 / pow(_h, 2);
                    vec_f(cnt) = func.laplace_f(gridtable[i][j]);
                }
                cnt++;
            }
        }
    }
    else if (bound == BoundCondiType::Neumann)
    {
        for (int i = 0; i <= _n; i++)
        {
            for (int j = 0; j <= _n; j++)
            {
                if (gridtable[i][j].insquare == -1) // corner
                {
                    Ah(cnt, cnt) = 1;
                    vec_f(cnt) = func(gridtable[i][j]);
                }
                else if (gridtable[i][j].insquare == 0) // inside
                {
                    Ah(cnt, cnt) = 4;
                    Ah(cnt, cnt + 1) = -1;
                    Ah(cnt, cnt - 1) = -1;
                    Ah(cnt, cnt + _n + 1) = -1;
                    Ah(cnt, cnt - _n - 1) = -1;
                    vec_f(cnt) = func.laplace_f(gridtable[i][j]) * pow(h, 2);
                }
                else if (gridtable[i][j].insquare == 1) // left
                {
                    Ah(cnt, cnt) = 4;
                    Ah(cnt, cnt + 1) = -2;
                    Ah(cnt, cnt + _n + 1) = -1;
                    Ah(cnt, cnt - _n - 1) = -1;
                    vec_f(cnt) = pow(h, 2) * func(gridtable[i][j]) - func.f_x(gridtable[i][j]) * 2.0 * h;
                }
                else if (gridtable[i][j].insquare == 2) // right
                {
                    Ah(cnt, cnt) = 4;
                    Ah(cnt, cnt - 1) = -2;
                    Ah(cnt, cnt + _n + 1) = -1;
                    Ah(cnt, cnt - _n - 1) = -1;
                    vec_f(cnt) = pow(h, 2) * func(gridtable[i][j]) + func.f_x(gridtable[i][j]) * 2.0 * h;
                }
                else if (gridtable[i][j].insquare == 3) // up
                {
                    Ah(cnt, cnt) = 4;
                    Ah(cnt, cnt - 1) = -1;
                    Ah(cnt, cnt + 1) = -1;
                    Ah(cnt, cnt - _n - 1) = -2;
                    vec_f(cnt) = pow(h, 2) * func(gridtable[i][j]) + func.f_y(gridtable[i][j]) * 2.0 * h;
                }
                else if (gridtable[i][j].insquare == 4) // down
                {
                    Ah(cnt, cnt) = 4;
                    Ah(cnt, cnt + 1) = -1;
                    Ah(cnt, cnt - 1) = -1;
                    Ah(cnt, cnt + _n + 1) = -2;
                    vec_f(cnt) = pow(h, 2) * func(gridtable[i][j]) - func.f_y(gridtable[i][j]) * 2.0 * h;
                }
                cnt++;
            }
        }
    }
    // cout << Ah << endl;
}

#endif